{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "import math\n",
    "from typing import Optional, List\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "GQA(group query attention)的核心是在多个注意力头中共享key value矩阵，从而减少key value矩阵的计算量。\n",
    "\n",
    "具体上，GQA将多个heads分成一组，一组内只计算一次key value矩阵，然后将结果广播到其他head中。当group_size=heads时，GQA被压缩成MQA，即所有heads共享一个key value矩阵。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from math import e\n",
    "import re\n",
    "from numpy import repeat\n",
    "\n",
    "\n",
    "class GroupQueryAttention(nn.Module):\n",
    "    def __init__(self, heads: int, d_model: int, group_size: int = 1, dropout_prob: float = 0.1, bias: bool = True):\n",
    "        super(GroupQueryAttention, self).__init__()\n",
    "        assert d_model % heads == 0\n",
    "        assert heads % group_size == 0\n",
    "\n",
    "        self.heads = heads\n",
    "        self.d_model = d_model\n",
    "        self.group_size = group_size\n",
    "        self.d_k = d_model//heads\n",
    "        self.num_group = heads//group_size\n",
    "\n",
    "        self.q_proj = nn.Linear(d_model, d_model, bias=bias)\n",
    "        self.k_proj = nn.Linear(d_model, self.d_k*self.num_group, bias=bias)\n",
    "        self.v_proj = nn.Linear(d_model, self.d_k*self.num_group, bias=bias)\n",
    "        self.output = nn.Linear(d_model, d_model)\n",
    "\n",
    "        self.softmax = nn.Softmax(dim=1)\n",
    "        self.dropout = nn.Dropout(dropout_prob)\n",
    "\n",
    "        self.scale = 1/math.sqrt(self.d_k)\n",
    "\n",
    "        self.attn = None\n",
    "\n",
    "    def get_score(self, query: torch.Tensor, key: torch.Tensor):\n",
    "        \"\"\"\n",
    "        Q shape:(seq_len_q, batch_size, num_heads, head_dim)\n",
    "        K shape:(seq_len_k, batch_size, num_heads//group_size, head_dim)\n",
    "        score shape:(seq_len_q, seq_len_k, batch_size, num_heads)\n",
    "        \"\"\"\n",
    "\n",
    "        return torch.einsum('qbhd,kbhd->qkbh', query, torch.repeat_interleave(key,repeats=self.group_size, dim=2))\n",
    "        \n",
    "    \n",
    "    def validate_mask(self, mask: torch.Tensor, query_shape: List[int], key_shape: List[int]):\n",
    "        \"\"\"校验mask\n",
    "\n",
    "        Args:\n",
    "            mask (torch.Tensor): 整体的mask，shape：(seq_len_q, seq_len_k, batch_size)\n",
    "            query_shape (List[int]): q的形状\n",
    "            key_shape (List[int]): k的形状\n",
    "        \"\"\"\n",
    "        assert mask.dim() >= 2\n",
    "        assert mask.shape[0] == 1 or mask.shape[0] == query_shape[0]\n",
    "        assert mask.shape[1] == key_shape[0]\n",
    "        if mask.dim() >= 3:\n",
    "            assert mask.shape[2] == 1 or mask.shape[2] == query_shape[1]\n",
    "        while mask.dim() < len(query_shape)+1:\n",
    "            mask.unsqueeze_(-1)\n",
    "        return mask\n",
    "    \n",
    "    \n",
    "    def forward(self, *, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: Optional[torch.Tensor] = None):\n",
    "        seq_len_k, batch_size, _ = key.shape\n",
    "        seq_len_q, batch_size, _ = query.shape\n",
    "\n",
    "        if mask is not None:\n",
    "            mask = self.validate_mask(mask, query.shape, key.shape)\n",
    "\n",
    "        query = self.q_proj(query).view(seq_len_k, batch_size, -1, self.d_k)\n",
    "        key = self.k_proj(key).view(seq_len_k, batch_size, -1, self.d_k)\n",
    "        value = self.v_proj(value).view(seq_len_k, batch_size, -1, self.d_k)\n",
    "\n",
    "        # (seq_len_q, seq_len_k, batch_size, num_heads)\n",
    "        scores = self.get_score(query, key)*self.scale\n",
    "        if mask is not None:\n",
    "            scores = scores.masked_fill(mask == 0, float('-inf'))\n",
    "\n",
    "        attn = self.softmax(scores)\n",
    "        attn = self.dropout(attn)\n",
    "        repeated_value = torch.repeat_interleave(value,repeats=self.group_size, dim=2)\n",
    "        \n",
    "        # einsum might produce discontigous result\n",
    "        x = torch.einsum(\"qkbh,kbhd->qbhd\", attn, repeated_value)\n",
    "        \n",
    "        self.attn = attn.detach()\n",
    "        x = x.reshape(seq_len_q, batch_size, -1)\n",
    "        return self.output(x)\n",
    "    \n",
    "if __name__ == '__main__':\n",
    "    heads, d_model, seq_len_q, batch, group_size = 8, 64, 5, 2, 4\n",
    "    mha = GroupQueryAttention(heads=heads, d_model=d_model, group_size=group_size)\n",
    "\n",
    "    q = torch.randn((seq_len_q, batch, d_model))\n",
    "    k = torch.randn((seq_len_q, batch, d_model))\n",
    "    v = torch.randn((seq_len_q, batch, d_model))\n",
    "    mask = torch.tril(torch.ones((seq_len_q, seq_len_q)))\n",
    "\n",
    "    output = mha(query=q, key=k, value=v, mask=mask)\n",
    "    assert output.shape==(seq_len_q, batch, d_model)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
